/*
 * @lc app=leetcode.cn id=86 lang=cpp
 *
 * [86] 分隔链表
 */

// @lc code=start
/**
 * Definition for singly-linked list.
 * struct ListNode {
 *     int val;
 *     ListNode *next;
 *     ListNode() : val(0), next(nullptr) {}
 *     ListNode(int x) : val(x), next(nullptr) {}
 *     ListNode(int x, ListNode *next) : val(x), next(next) {}
 * };
 */
#include <bits/stdc++.h>
using namespace std;
class Solution
{
public:
    ListNode *partition(ListNode *head, int x)
    {
        if (!head)
            return head;
        vector<ListNode *> less_nodes, greater_nodes;
        while (head)
        {
            if (head->val < x)
            {
                less_nodes.push_back(head);
            }
            else
            {
                greater_nodes.push_back(head);
            }
            head = head->next;
        }

        if (less_nodes.size())
        {
            head = less_nodes[0];
        }
        else
            head = greater_nodes[0];

        auto p = head;
        for (int i = 1; i < less_nodes.size(); i++)
        {
            p->next = less_nodes[i];
            p = less_nodes[i];
        }
        p->next = nullptr;
        int i = less_nodes.size() ? 0 : 1;
        while (i < greater_nodes.size())
        {
            p->next = greater_nodes[i];
            p = greater_nodes[i];
            i++;
        }
        p->next = nullptr;
        return head;
    }
};
// @lc code=end
